package simpletrace

import (
	"context"
	"fmt"
	"hash/crc32"
	"time"
)

type DefaultExporter struct {
	taskChan          []chan *Span
	taskSize          int
	maxTaskSpanSize   int
	ticker            *time.Ticker
	doneChan          chan struct{}
	setTraceIDFn      SetTraceIDFn
	setTraceSpansFn   SetTraceSpansFn
	onTraceExportedFn OnTraceExportedFn
}

type SetTraceIDFn func(ctx context.Context, traceID, spanID string) error

type SetTraceSpansFn func(ctx context.Context, span []*Span) error

type OnTraceExportedFn func(ctx context.Context, traceID string)

type DefaultExporterOptionFn func(*DefaultExporter)

func WithDefTaskSize(size int) DefaultExporterOptionFn {
	return func(de *DefaultExporter) {
		if size > 0 {
			de.taskSize = size
		}
	}
}

func WithDefMaxTaskSpanSize(size int) DefaultExporterOptionFn {
	return func(de *DefaultExporter) {
		if size > 0 {
			de.maxTaskSpanSize = size
		}
	}
}

func WithDefSetTraceIDFn(fn SetTraceIDFn) DefaultExporterOptionFn {
	return func(de *DefaultExporter) {
		de.setTraceIDFn = fn
	}
}

func WithDefSetTraceSpans(fn SetTraceSpansFn) DefaultExporterOptionFn {
	return func(de *DefaultExporter) {
		de.setTraceSpansFn = fn
	}
}

func WithDefOnTraceExportedFn(fn OnTraceExportedFn) DefaultExporterOptionFn {
	return func(de *DefaultExporter) {
		de.onTraceExportedFn = fn
	}
}

func NewDefaultExporter(opts ...DefaultExporterOptionFn) Exporter {
	exporter := &DefaultExporter{
		ticker:   time.NewTicker(time.Second * 10),
		doneChan: make(chan struct{}),
	}
	for _, opt := range opts {
		opt(exporter)
	}
	if exporter.taskSize <= 0 {
		exporter.taskSize = 4
	}
	if exporter.maxTaskSpanSize <= 0 {
		exporter.maxTaskSpanSize = 5
	}
	exporter.taskChan = make([]chan *Span, exporter.taskSize)
	for i := range exporter.taskSize {
		exporter.taskChan[i] = make(chan *Span, exporter.maxTaskSpanSize*4)
	}
	return exporter
}

func (re *DefaultExporter) Export(ctx context.Context, span *Span) error {
	if span == nil || span.TraceID == "" || span.SpanID == "" {
		logger.WarnCtx(ctx, "[strace] span missing TraceID or SpanID, skipping export")
		return nil
	}
	index := crc32.ChecksumIEEE([]byte(span.TraceID)) % uint32(re.taskSize)
	if index >= uint32(len(re.taskChan)) {
		return fmt.Errorf("[strace] invalid task channel index: %d", index)
	}
	re.taskChan[index] <- span
	return nil
}

func (re *DefaultExporter) Start(ctx context.Context) {
	for i := range re.taskSize {
		task := &DefaultExporterTask{
			Spans:             make([]*Span, 0, re.maxTaskSpanSize),
			ticker:            re.ticker,
			doneChan:          re.doneChan,
			defaultSize:       re.maxTaskSpanSize,
			taskChan:          re.taskChan[i],
			traces:            make(map[string]bool),
			setTraceIDFn:      re.setTraceIDFn,
			setTraceSpansFn:   re.setTraceSpansFn,
			onTraceExportedFn: re.onTraceExportedFn,
		}
		go task.run(ctx)
	}
}

func (re *DefaultExporter) Close() {
	re.ticker.Stop()
	close(re.doneChan)
	for _, taskChan := range re.taskChan {
		close(taskChan) // Close each task channel to signal completion
	}
	logger.Info("DefaultExporter closed")
}

type DefaultExporterTask struct {
	Spans       []*Span
	ticker      *time.Ticker
	doneChan    chan struct{}
	defaultSize int
	taskChan    chan *Span
	traces      map[string]bool

	setTraceIDFn      SetTraceIDFn
	setTraceSpansFn   SetTraceSpansFn
	onTraceExportedFn OnTraceExportedFn
}

func (re *DefaultExporterTask) run(ctx context.Context) {
	for {
		defer func() {
			if r := recover(); r != nil {
				logger.ErrorCtx(ctx, "[strace] recovered from panic: %v", r)
			}
			re.ticker.Stop()
			close(re.doneChan)
		}()
		select {
		case span := <-re.taskChan:

			if span == nil {
				continue // Skip nil spans
			}
			if span.TraceID == "" || span.SpanID == "" {
				logger.WarnCtx(ctx, "[strace] span missing TraceID or SpanID, skipping export")
				continue
			}
			if _, exists := re.traces[span.TraceID]; !exists {
				re.traces[span.TraceID] = true
				err := re.SetTraceID(ctx, span.TraceID, span.SpanID)
				if err != nil {
					logger.ErrorCtx(ctx, "[strace] failed to set trace ID: %v", err)
					continue
				}
				re.OnTraceExported(ctx, span.TraceID)
			}

			re.Spans = append(re.Spans, span)
			if len(re.Spans) >= re.defaultSize {
				// Process the batch of spans
				re.flush(ctx)
				re.Spans = nil // Reset for the next batch
			}
		case <-re.ticker.C:
			if len(re.Spans) > 0 {
				re.flush(ctx)
				re.Spans = nil // Reset after processing
			}
		case <-re.doneChan:
			return // Exit the goroutine when done
		}
	}
}

func (re *DefaultExporterTask) flush(ctx context.Context) {
	if len(re.Spans) == 0 {
		return
	}

	err := re.SetTraceSpans(ctx, re.Spans)
	if err != nil {
		logger.ErrorCtx(ctx, "[strace] failed to set trace span infos: %v", err)
		return
	}

	logger.InfoCtx(ctx, "[strace] successfully exported %d spans", len(re.Spans))
}

func (re *DefaultExporterTask) SetTraceID(ctx context.Context, traceID, spanID string) error {
	if re.setTraceIDFn == nil {
		return nil
	}
	return re.setTraceIDFn(ctx, traceID, spanID)
}

func (re *DefaultExporterTask) SetTraceSpans(ctx context.Context, span []*Span) error {
	if re.setTraceSpansFn == nil {
		return nil
	}
	return re.setTraceSpansFn(ctx, span)
}

func (re *DefaultExporterTask) OnTraceExported(ctx context.Context, traceID string) {
	if re.onTraceExportedFn == nil {
		return
	}
	re.onTraceExportedFn(ctx, traceID)
}
